//MIT License
//Copyright (c) 2024 nuaazgg, Nanjing University of Aeronautics and Astronautics(NUAA)

#include <filesystem/filesystem.h>
#include <icg/azure_kinect_camera.h>
#include <icg/basic_depth_renderer.h>
#include <icg/body.h>
#include <icg/common.h>
#include <icg/depth_modality.h>
#include <icg/depth_model.h>
#include <icg/normal_viewer.h>
#include <icg/region_modality.h>
#include <icg/region_model.h>
#include <icg/renderer_geometry.h>
#include <icg/static_detector.h>
#include <icg/tracker.h>

#include <Eigen/Geometry>
#include <memory>
#include <string>
#include <ppfmatch.h>
#include <libwebsockets.h>
#include <thread>

Pose_3D pose;
// WebSocket连接状态
enum ConnectionState {
  CONN_DISCONNECTED,
  CONN_CONNECTED,
};

// WebSocket上下文
struct WebsocketContext {
  lws_context* context;
  ConnectionState state;
};

// WebSocket事件处理函数
int callback_websocket(struct lws* wsi, enum lws_callback_reasons reason, void* user, void* in, size_t len) {
  switch (reason) {
      case LWS_CALLBACK_ESTABLISHED:
        cout << "WebSocket connection established" << endl;
        break;

      case LWS_CALLBACK_SERVER_WRITEABLE:
        break;

      case LWS_CALLBACK_RECEIVE:
        len = sizeof(pose)
        lws_write(wsi, (unsigned char*)pose, len, LWS_WRITE_TEXT);
        break;

      default:
        break;
  }

  return 0;
}

// WebSocket服务器线程函数
void websocket_server_thread(WebsocketContext* ctx) {
  // 运行WebSocket事件循环
  while (true) {
      lws_service(ctx->context, 50);
  }
}

int main(int argc, char *argv[]) {
  if (argc < 3) {
    std::cerr << "Not enough arguments: Provide directory and body_names";
    return 0;
  }
  const std::filesystem::path directory{argv[1]};
  std::vector<std::string> body_names;
  for (int i = 2; i < argc; ++i) {
    body_names.push_back(std::string{argv[i]});
  }

  constexpr bool kUseDepthViewer = true;
  constexpr bool kMeasureOcclusions = true;
  constexpr bool kModelOcclusions = false;
  constexpr bool kVisualizePoseResult = false;
  constexpr bool kSaveImages = false;
  const std::filesystem::path save_directory{""};

  // Set up tracker and renderer geometry
  auto tracker_ptr{std::make_shared<icg::Tracker>("tracker")};

  // Set up cameras
  auto color_camera_ptr{
      std::make_shared<icg::AzureKinectColorCamera>("azure_kinect_color")};
  auto depth_camera_ptr{
      std::make_shared<icg::AzureKinectDepthCamera>("azure_kinect_depth")};

  for (const auto body_name : body_names) {
    // Set up body
    std::filesystem::path metafile_path{directory / (body_name + ".yaml")};
    auto body_ptr{std::make_shared<icg::Body>(body_name, metafile_path)};

    // Set up detector
    std::filesystem::path detector_path{directory /
                                        (body_name + "_detector.yaml")};
    auto detector_ptr{std::make_shared<icg::StaticDetector>(
        body_name + "_detector", detector_path, body_ptr)};
    tracker_ptr->AddDetector(detector_ptr);

    // Set up models
    auto region_model_ptr{std::make_shared<icg::RegionModel>(
        body_name + "_region_model", body_ptr,
        directory / (body_name + "_region_model.bin"))};
    auto depth_model_ptr{std::make_shared<icg::DepthModel>(
        body_name + "_depth_model", body_ptr,
        directory / (body_name + "_depth_model.bin"))};

    // Set up modalities
    auto region_modality_ptr{std::make_shared<icg::RegionModality>(
        body_name + "_region_modality", body_ptr, color_camera_ptr,
        region_model_ptr)};
    auto depth_modality_ptr{std::make_shared<icg::DepthModality>(
        body_name + "_depth_modality", body_ptr, depth_camera_ptr,
        depth_model_ptr)};
    if (kMeasureOcclusions) {
      region_modality_ptr->MeasureOcclusions(depth_camera_ptr);
      depth_modality_ptr->MeasureOcclusions();
    }
    if (kModelOcclusions) {
      region_modality_ptr->ModelOcclusions(color_depth_renderer_ptr);
      depth_modality_ptr->ModelOcclusions(depth_depth_renderer_ptr);
    }

    // Set up optimizer
    auto body1_optimizer_ptr{
        std::make_shared<icg::Optimizer>(body_name + "_optimizer")};
    body1_optimizer_ptr->AddModality(region_modality_ptr);
    body1_optimizer_ptr->AddModality(depth_modality_ptr);
    tracker_ptr->AddOptimizer(body1_optimizer_ptr);
  }


  //setup init
  match = new PPFmatch();
  double samplingStep = 0.05;
  double distanceStep = 0.04;
  int numAngles = 15;
  std::filesystem::path modelFileName{directory /
                                    (body_name + "pcd")};
  matcher->setParameter(samplingStep, distanceStep, numAngles);
  matcher->train(modelFileName);

  vector<Pose_3D, Eigen::aligned_allocator<Pose_3D>> result;
  double scale = 0.3;
  matcher->loadScenePointCloudFile(sceneFileName);
  matcher->match(sceneFileName, result, scale, distanceStep);

  // WebSocket
  const int port = 8080;

  WebsocketContext ctx;
  ctx.state = CONN_DISCONNECTED;

  struct lws_context_creation_info info;
  memset(&info, 0, sizeof(info));
  info.port = port;
  info.iface = nullptr;
  info.protocols = lws_protocols;
  info.extensions = nullptr;
  info.gid = -1;
  info.uid = -1;
  info.options = 0;
  ctx.context = lws_create_context(&info);
  if (ctx.context == nullptr) {
      cerr << "Failed to create WebSocket context" << endl;
      return 1;
  }

  cout << "WebSocket server started on port " << port << endl;

  // 创建WebSocket服务器线程
  thread server_thread(websocket_server_thread, &ctx);

  // 等待线程结束
  server_thread.join();


  // Start tracking
  if (!tracker_ptr->SetUp()) return 0;
  tracker_ptr->Initlocation(result[0],&pose);
  if (!tracker_ptr->RunTrackerProcess(true, false)) return 0;
  lws_context_destroy(ctx.context);
  return 0;
}

